"""Shapefile数据导入工具

用于将Shapefile数据导入到PostGIS数据库中
"""

import logging
import os
import tempfile
import traceback
import zipfile
from pathlib import Path
from typing import Optional, Dict

import geopandas as gpd
from sqlalchemy import create_engine, text

from hydropulse.config.settings import get_database_url

logger = logging.getLogger(__name__)


class ShapefileImporter:
    """Shapefile导入器"""

    def __init__(self):
        """初始化导入器"""
        self.engine = create_engine(get_database_url())

    def extract_shapefile_from_zip(self, zip_path: str, temp_dir: str) -> Optional[str]:
        """从ZIP文件中提取Shapefile
        
        Args:
            zip_path: ZIP文件路径
            temp_dir: 临时目录路径
            
        Returns:
            Shapefile路径（.shp文件）
        """
        try:
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(temp_dir)

            # 查找.shp文件
            for root, dirs, files in os.walk(temp_dir):
                for file in files:
                    if file.endswith('.shp'):
                        return os.path.join(root, file)

            logger.error(f"在ZIP文件 {zip_path} 中未找到.shp文件")
            return None

        except Exception as e:
            logger.error(f"提取ZIP文件失败: {e}")
            return None

    def _read_shapefile_with_encoding(self, shp_path: str):
        """尝试使用不同编码读取Shapefile
        
        Args:
            shp_path: Shapefile路径
            
        Returns:
            GeoDataFrame对象
        """
        # 常见的中文编码列表
        encodings = ['utf-8', 'gbk', 'gb2312', 'gb18030', 'cp936']
        
        for encoding in encodings:
            try:
                logger.info(f"尝试使用 {encoding} 编码读取文件")
                gdf = gpd.read_file(shp_path, encoding=encoding)
                logger.info(f"成功使用 {encoding} 编码读取文件")
                return gdf
            except UnicodeDecodeError as e:
                logger.warning(f"使用 {encoding} 编码读取失败: {e}")
                continue
            except Exception as e:
                logger.error(f"使用 {encoding} 编码读取时发生其他错误: {e}")
                continue
        
        raise Exception("所有编码尝试都失败了")

    def import_provinces(self, zip_path: str) -> bool:
        """导入省界数据
        
        Args:
            zip_path: 省界ZIP文件路径
            
        Returns:
            是否导入成功
        """
        logger.info(f"开始导入省界数据: {zip_path}")

        with tempfile.TemporaryDirectory() as temp_dir:
            # 提取Shapefile
            shp_path = self.extract_shapefile_from_zip(zip_path, temp_dir)
            if not shp_path:
                return False

            try:
                # 读取Shapefile（尝试多种编码）
                gdf = self._read_shapefile_with_encoding(shp_path)
                logger.info(f"读取到 {len(gdf)} 个省级行政区")

                # 确保几何列名为geometry
                if 'geometry' not in gdf.columns:
                    geom_col = gdf.geometry.name
                    gdf = gdf.rename(columns={geom_col: 'geometry'})

                # 设置坐标系为WGS84 (EPSG:4326)
                if gdf.crs is None:
                    gdf.set_crs('EPSG:4326', inplace=True)
                elif gdf.crs != 'EPSG:4326':
                    gdf = gdf.to_crs('EPSG:4326')
                logger.info(gdf.columns)

                # 从数据库查询已有的省份列表
                with self.engine.connect() as conn:
                    result = conn.execute(
                        text("SELECT name FROM spatial_regions WHERE region_type = '省'")
                    )
                    target_provinces = [row[0] for row in result.fetchall()]
                    logger.info(f"数据库中已有省份: {target_provinces}")
                
                # 准备数据
                processed_data = []
                for idx, row in gdf.iterrows():
                    # 根据实际字段获取省名
                    name = None
                    if '省' in row and row['省'] is not None:
                        name = str(row['省']).strip()
                    elif 'Name' in row and row['Name'] is not None:
                        name = str(row['Name']).strip()
                    elif 'name' in row and row['name'] is not None:
                        name = str(row['name']).strip()
                    
                    if not name:
                        name = f"省级行政区_{idx + 1}"

                    # 过滤：只处理目标省份
                    if not any(target in name for target in target_provinces):
                        continue

                    # 获取省代码
                    province_code = None
                    if '省代码' in row and row['省代码'] is not None:
                        province_code = str(row['省代码']).strip()
                    elif 'code' in row and row['code'] is not None:
                        province_code = str(row['code']).strip()
                    elif 'adcode' in row and row['adcode'] is not None:
                        province_code = str(row['adcode']).strip()

                    processed_data.append({
                        'name': name,
                        'description': f'{name}省级行政区域',
                        'geometry': row['geometry'],
                        'region_type': '省',
                        'level': 1,
                        'code': province_code
                    })

                # 创建GeoDataFrame
                import_gdf = gpd.GeoDataFrame(processed_data, crs='EPSG:4326')

                # 基于省份名称去重导入
                with self.engine.connect() as conn:
                    for _, row in import_gdf.iterrows():
                        # 检查是否已存在相同名称的省份
                        existing = conn.execute(
                            text("SELECT COUNT(*) FROM spatial_regions WHERE name = :name AND region_type = '省'"),
                            {'name': row['name']}
                        ).scalar()
                        
                        if existing == 0:
                             # 不存在则插入
                             # 处理几何数据，移除Z维度
                             geom = row['geometry']
                             if geom.has_z:
                                 # 移除Z维度，只保留X,Y坐标
                                 from shapely.ops import transform
                                 geom = transform(lambda x, y, z=None: (x, y), geom)
                             
                             conn.execute(
                                 text("""
                                 INSERT INTO spatial_regions (name, description, geometry, region_type, level, code)
                                 VALUES (:name, :description, ST_GeomFromText(:geometry, 4326), :region_type, :level, :code)
                                 """),
                                 {
                                     'name': row['name'],
                                     'description': row['description'],
                                     'geometry': geom.wkt,
                                     'region_type': row['region_type'],
                                     'level': row['level'],
                                     'code': row['code']
                                 }
                             )
                             logger.info(f"插入新省份: {row['name']}")
                        else:
                            logger.info(f"省份已存在，跳过: {row['name']}")
                    
                    conn.commit()

                logger.info(f"成功导入 {len(import_gdf)} 个省级行政区")
                return True

            except Exception as e:
                traceback.print_exc()
                logger.error(f"导入省界数据失败: {e}")
                return False

    def import_cities(self, zip_path: str) -> bool:
        """导入市界数据
        
        Args:
            zip_path: 市界ZIP文件路径
            
        Returns:
            是否导入成功
        """
        logger.info(f"开始导入市界数据: {zip_path}")

        with tempfile.TemporaryDirectory() as temp_dir:
            # 提取Shapefile
            shp_path = self.extract_shapefile_from_zip(zip_path, temp_dir)
            if not shp_path:
                return False

            try:
                # 读取Shapefile（尝试多种编码）
                gdf = self._read_shapefile_with_encoding(shp_path)
                logger.info(f"读取到 {len(gdf)} 个市级行政区")

                # 确保几何列名为geometry
                if 'geometry' not in gdf.columns:
                    geom_col = gdf.geometry.name
                    gdf = gdf.rename(columns={geom_col: 'geometry'})

                # 设置坐标系为WGS84 (EPSG:4326)
                if gdf.crs is None:
                    gdf.set_crs('EPSG:4326', inplace=True)
                elif gdf.crs != 'EPSG:4326':
                    gdf = gdf.to_crs('EPSG:4326')
                logger.info(gdf.columns)

                # 从数据库查询已有的省份列表
                with self.engine.connect() as conn:
                    result = conn.execute(
                        text("SELECT name FROM spatial_regions WHERE region_type = '省'")
                    )
                    target_provinces = [row[0] for row in result.fetchall()]
                    logger.info(f"数据库中已有省份: {target_provinces}")
                
                # 准备数据
                processed_data = []
                for idx, row in gdf.iterrows():
                    # 根据实际字段获取市名
                    name = None
                    if 'Name' in row and row['Name'] is not None:
                        name = str(row['Name']).strip()
                    elif 'name' in row and row['name'] is not None:
                        name = str(row['name']).strip()
                    
                    if not name:
                        name = f"市级行政区_{idx + 1}"

                    # 获取省名
                    province_name = None
                    if '省' in row and row['省'] is not None:
                        province_name = str(row['省']).strip()

                    # 过滤：只处理目标省份下的市
                    if province_name and not any(target in province_name for target in target_provinces):
                        continue

                    # 获取省代码
                    province_code = None
                    if '省代码' in row and row['省代码'] is not None:
                        province_code = str(row['省代码']).strip()

                    # 获取市代码
                    city_code = None
                    if 'code' in row and row['code'] is not None:
                        city_code = str(row['code']).strip()

                    # 计算面积（平方米）
                    area = 0
                    try:
                        # 如果有Shape_Area字段，直接使用
                        if 'Shape_Area' in row and row['Shape_Area'] is not None:
                            area = int(float(row['Shape_Area']))
                        else:
                            # 转换到等面积投影计算面积
                            geom_proj = gpd.GeoSeries([row['geometry']], crs='EPSG:4326').to_crs('EPSG:3857')
                            area = int(geom_proj.area.iloc[0])
                    except:
                        area = 0

                    # 确定类别
                    category = '地级市'
                    if '市' in name:
                        if any(x in name for x in ['北京', '上海', '天津', '重庆']):
                            category = '直辖市'
                        elif '副省级' in str(row.get('grade', '')) or name in ['大连市', '青岛市', '宁波市', '厦门市', '深圳市']:
                            category = '副省级市'

                    processed_data.append({
                        'name': name,
                        'description': f'{name}行政区域',
                        'geometry': row['geometry'],
                        'category': category,
                        'area': area,
                        'province_name': province_name,
                        'province_code': province_code,
                        'city_code': city_code
                    })

                # 创建GeoDataFrame
                import_gdf = gpd.GeoDataFrame(processed_data, crs='EPSG:4326')

                # 基于城市名称去重导入
                with self.engine.connect() as conn:
                    for _, row in import_gdf.iterrows():
                        # 检查是否已存在相同名称的城市
                        existing = conn.execute(
                            text("SELECT COUNT(*) FROM spatial_polygons WHERE name = :name AND category IN ('直辖市', '副省级市', '地级市')"),
                            {'name': row['name']}
                        ).scalar()
                        
                        if existing == 0:
                            # 不存在则插入
                            # 处理几何数据，移除Z维度
                            geom = row['geometry']
                            if geom.has_z:
                                # 移除Z维度，只保留X,Y坐标
                                from shapely.ops import transform
                                geom = transform(lambda x, y, z=None: (x, y), geom)
                            
                            conn.execute(
                                text("""
                                INSERT INTO spatial_polygons (name, description, geometry, category, area, province_name, province_code, city_code)
                                VALUES (:name, :description, ST_GeomFromText(:geometry, 4326), :category, :area, :province_name, :province_code, :city_code)
                                """),
                                {
                                    'name': row['name'],
                                    'description': row['description'],
                                    'geometry': geom.wkt,
                                    'category': row['category'],
                                    'area': row['area'],
                                    'province_name': row['province_name'],
                                    'province_code': row['province_code'],
                                    'city_code': row['city_code']
                                }
                            )
                            logger.info(f"插入新城市: {row['name']}")
                        else:
                            logger.info(f"城市已存在，跳过: {row['name']}")
                    
                    conn.commit()

                logger.info(f"成功导入 {len(import_gdf)} 个市级行政区")
                return True

            except Exception as e:
                traceback.print_exc()
                logger.error(f"导入市界数据失败: {e}")
                return False

    def import_all_data(self, province_zip: str, city_zip: str) -> Dict[str, bool]:
        """导入所有数据
        
        Args:
            province_zip: 省界ZIP文件路径
            city_zip: 市界ZIP文件路径
            
        Returns:
            导入结果字典
        """
        results = {
            'provinces': False,
            'cities': False
        }

        # 导入省界数据
        if os.path.exists(province_zip):
            results['provinces'] = self.import_provinces(province_zip)
        else:
            logger.error(f"省界文件不存在: {province_zip}")

        # 导入市界数据
        if os.path.exists(city_zip):
            results['cities'] = self.import_cities(city_zip)
        else:
            logger.error(f"市界文件不存在: {city_zip}")

        return results

    def update_region_geometry_from_shapefile(self, zip_path: str, region_id: int) -> bool:
        """从Shapefile更新指定区域的几何数据
        
        Args:
            zip_path: Shapefile ZIP文件路径
            region_id: 要更新的区域ID
            
        Returns:
            是否更新成功
        """
        logger.info(f"开始更新区域ID {region_id} 的几何数据: {zip_path}")

        with tempfile.TemporaryDirectory() as temp_dir:
            # 提取Shapefile
            shp_path = self.extract_shapefile_from_zip(zip_path, temp_dir)
            if not shp_path:
                return False

            try:
                # 读取Shapefile（尝试多种编码）
                gdf = self._read_shapefile_with_encoding(shp_path)
                logger.info(f"读取到 {len(gdf)} 个几何要素")

                # 确保几何列名为geometry
                if 'geometry' not in gdf.columns:
                    geom_col = gdf.geometry.name
                    gdf = gdf.rename(columns={geom_col: 'geometry'})

                # 设置坐标系为WGS84 (EPSG:4326)
                if gdf.crs is None:
                    gdf.set_crs('EPSG:4326', inplace=True)
                elif gdf.crs != 'EPSG:4326':
                    gdf = gdf.to_crs('EPSG:4326')

                # 检查区域是否存在
                with self.engine.connect() as conn:
                    existing = conn.execute(
                        text("SELECT name FROM spatial_regions WHERE id = :id"),
                        {'id': region_id}
                    ).fetchone()
                    
                    if not existing:
                        logger.error(f"区域ID {region_id} 不存在")
                        return False
                    
                    region_name = existing[0]
                    logger.info(f"找到区域: {region_name}")

                    # 选择面积最大的几何要素
                    if len(gdf) == 0:
                        logger.error("Shapefile中没有几何要素")
                        return False
                    
                    # 计算每个几何要素的面积并选择最大的
                    gdf['area'] = gdf['geometry'].area
                    max_area_idx = gdf['area'].idxmax()
                    geometry = gdf.loc[max_area_idx, 'geometry']
                    
                    logger.info(f"选择了面积最大的几何要素，面积: {gdf.loc[max_area_idx, 'area']:.6f}")
                    
                    # 处理几何数据，移除Z维度
                    if geometry.has_z:
                        from shapely.ops import transform
                        geometry = transform(lambda x, y, z=None: (x, y), geometry)
                    
                    # 更新几何数据
                    conn.execute(
                        text("""
                        UPDATE spatial_regions 
                        SET geometry = ST_GeomFromText(:geometry, 4326)
                        WHERE id = :id
                        """),
                        {
                            'geometry': geometry.wkt,
                            'id': region_id
                        }
                    )
                    conn.commit()
                    
                    logger.info(f"成功更新区域ID {region_id} ({region_name}) 的几何数据")
                    return True

            except Exception as e:
                traceback.print_exc()
                logger.error(f"更新几何数据失败: {e}")
                return False


def main():
    """主函数 - 执行数据导入"""
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

    # 文件路径
    base_dir = Path(__file__).parent.parent.parent.parent
    province_zip = base_dir / "data/全国_省界.zip"
    city_zip = base_dir / "data/全国_市界.zip"

    # 创建导入器
    importer = ShapefileImporter()

    # 执行导入
    logger.info("开始导入Shapefile数据...")
    results = importer.import_all_data(str(province_zip), str(city_zip))

    # 输出结果
    logger.info("导入结果:")
    logger.info(f"  省界数据: {'成功' if results['provinces'] else '失败'}")
    logger.info(f"  市界数据: {'成功' if results['cities'] else '失败'}")

    if all(results.values()):
        logger.info("所有数据导入成功！")
    else:
        logger.error("部分数据导入失败，请检查日志")


if __name__ == "__main__":
    main()
